from concurrent.futures import ThreadPoolExecutor

import pyarrow.compute as pc
import vortex as vx


def sum_column(path: str, column: str) -> int | float:
    reader = vx.open(path).to_arrow([column], batch_size=64_000)
    total = 0

    for batch in reader:
        value = pc.sum(batch.column(column)).as_py()
        if value is not None:
            total += value

    return total


columns = ["tip_amount", "fare_amount", "total_amount"]
with ThreadPoolExecutor(max_workers=len(columns)) as threads:
    totals = list(threads.map(lambda column: sum_column("example.vortex", column), columns))

import vortex as vx
import vortex.runtime as vxrt

previous_workers = vxrt.worker_count()
vxrt.set_worker_threads_to_available_parallelism()

try:
    reader = vx.open("example.vortex").to_arrow(batch_size=64_000)
    table = reader.read_all()
finally:
    vxrt.set_worker_threads(previous_workers)